//
// Created by benxb on 2021/12/12.
//


#include <index/linear_index.h>
#include <common/exception.h>
#include <common/util/file_util.h>
#include <algorithm>
#include <dirent.h>
#include <common/logger.h>
#include <index/linear_index_segment_reader.h>

namespace benxdb {

LinearIndex::~LinearIndex() {
    FlushIndex();
    index_io_.close();
}


LinearIndex::LinearIndex(const std::string &table_name, const Column &column) : column_(column) {
    switch (column_.GetType()) {
        case FLOAT:
            key_type_ = STYPE_FLOAT;
            break;
        case INTEGER:
            key_type_ = STYPE_INTEGER;
            break;
        case BIGINT:
            key_type_ = STYPE_BIGINT;
            break;
        case DECIMAL:
            key_type_ = STYPE_DECIMAL;
            break;
        case TIMESTAMP:
            key_type_ = STYPE_TIMESTAMP;
            break;
    };
    index_file_ = GlobalConfig::DatasetDir(table_name) + "/" + column.GetName() + INDEX_FILE_SUFFIX;
    index_data_dir_ = GlobalConfig::DatasetDir(table_name) + "/" + column.GetName() + INDEX_DATA_DIR_SUFFIX;
    InitIndexIO();
    InitIndexDataDIR();
    if (indexes_.empty()) {
        // create first index segment
        segments_.insert({0, std::make_shared<LinearIndexSegment>(index_data_dir_, 0, key_type_)});
        LinearIndexPair first_index;
        first_index.key = LinearIndexKey::GetMinValueOfType(key_type_);
        first_index.seg_id = 0;
        indexes_.push_back(first_index);
        FlushIndex();
    }
    page_ = std::make_shared<LinearIndexPage>();
}

/**
 * @brief generate index entry
 * 
 * @param data point cloud attrributes
 * @param size data size
 * @param begin_id first id of these data
 */
void LinearIndex::GenerateIndex(const char *data, int size, record_id_t begin_id) {
    std::vector<LinearIndexLeafPair> new_pairs;
    int offset = 0;
    record_id_t record_id = begin_id;
    // first generate index pairs
    while (offset < size) {
        new_pairs.emplace_back(LinearIndexLeafPair(key_type_, data + offset, record_id));
        offset += LinearIndexKey::GetSizeOfType(key_type_);
        record_id++;
    }
    GenerateIndex(new_pairs, 0, new_pairs.size() - 1, 0, indexes_.size() - 1);
    HandleSplitSegments();
}

void LinearIndex::InitIndexIO() {
    index_io_.open(index_file_, std::ios::binary | std::ios::in | std::ios::out);
    int file_size = FileUtil::GetFileSize(index_file_);
    // directory or file does not exist
    if (!index_io_.is_open()) {
        index_io_.clear();
        // create a new file
        index_io_.open(index_file_, std::ios::binary | std::ios::trunc | std::ios::out);
        index_io_.close();
        // reopen with original mode
        index_io_.open(index_file_, std::ios::binary | std::ios::in | std::ios::out);
        if (!index_io_.is_open()) {
            throw Exception("can't open index file");
        }
    }
    if (file_size > 0) {
        // Read index
        LOG_DEBUG("index file not empty.");
        while (!index_io_.eof()) {
            LinearIndexPair pair{};
            index_io_.read(reinterpret_cast<char *>(&pair), sizeof(pair));
            indexes_.push_back(pair);
        }
    }
}

void LinearIndex::InitIndexDataDIR() {
    if (!opendir(index_data_dir_.c_str())) {
        mkdir(index_data_dir_.c_str(), 0775);
    }
}

std::shared_ptr<LinearIndexSegment> LinearIndex::GetIndexSegment(int segment_id) {
    if (segments_.find(segment_id) == segments_.end()) {
        segments_[segment_id] = std::make_shared<LinearIndexSegment>(index_data_dir_, segment_id, key_type_);
    }
    return segments_[segment_id];
}

/**
 * @brief generate index and append into index files recursively
 * 
 * @param columns index pairs
 * @param cposL begin position of index pairs
 * @param cposH end position of index pairs
 * @param iposL begin index of index segement
 * @param iposH end index of index segment
 */
void LinearIndex::GenerateIndex(std::vector<LinearIndexLeafPair> &columns, int cposL, int cposH, int iposL, int iposH) {
    if (cposL > cposH || iposL > iposH) {
        return;
    }
    int iposMid = (iposL + iposH) / 2;
    LinearIndexPair &indexMid = indexes_[iposMid];
    LinearIndexKey key = indexMid.key;
    int cposMid = 0;
    if (iposMid > 0) {
        auto second_begin = std::stable_partition(columns.begin() + cposL, columns.begin() + cposH,
                                                  [key](const auto &col) {
                                                      return col.key.CompareLessThan(key);
                                                  });
        cposMid = second_begin - columns.begin();
    }
    if (iposL == iposH) {
        AppendIndex(iposL - 1, columns, cposL, cposMid - 1);
        AppendIndex(iposL, columns, cposMid, cposH);
    } else {
        GenerateIndex(columns, cposL, cposMid - 1, iposL, iposMid - 1);
        GenerateIndex(columns, cposMid, cposH, iposMid + 1, iposH);
    }

}

void LinearIndex::AppendIndex(int seg_index, std::vector<LinearIndexLeafPair> &column, int cposL, int cposH) {
    if (seg_index < 0 || seg_index >= indexes_.size()) {
        return;
    }
    if (cposL > cposH) {
        return;
    }
    std::shared_ptr<LinearIndexSegment> segment = GetIndexSegment(indexes_[seg_index].seg_id);
    segment->AppendManyIndexPairs(column, cposL, cposH);
    if (segment->GetPageCount() > SEGMENT_SPLIT_PAGE_COUNT) {
        need_split_.push_back(seg_index);
    }
}

void LinearIndex::FlushIndex() {
    index_io_.seekp(0);
    for (auto &pair : indexes_) {
        index_io_.write(reinterpret_cast<char *>(&pair), sizeof(pair));
    }
}

/**
 * @brief split index segment file
 * 
 * @param seg_index segment index, also index file name
 */
void LinearIndex::SplitSegment(int seg_index) {
    LinearIndexPair &pair = indexes_[seg_index];
    index_seg_id_t seg_id = pair.seg_id;
    std::shared_ptr<LinearIndexSegment> seg = GetIndexSegment(seg_id);
    const LinearIndexKey &split_key = seg->GetApproximateMid();
    std::shared_ptr<LinearIndexSegment> seg_0 = seg->CreateTempSegment(0);
    std::shared_ptr<LinearIndexSegment> seg_1 = seg->CreateTempSegment(1);
    LinearIndexSegmentReader reader(seg);
    while (reader.HasMoreIndex()) {
        auto leaf_pair = reader.GetNextIndexPair();
        if (leaf_pair.key.CompareLessThan(split_key)) {
            seg_0->AppendSingleIndexWithoutFlush(leaf_pair);
        } else {
            seg_1->AppendSingleIndexWithoutFlush(leaf_pair);
        }
    }
    seg_0->Flush();
    seg_1->Flush();
    seg->Destroy();
    seg_0->RenameSegment(seg_id);
    index_seg_id_t next_seg_id = GetNextSegmentId();
    seg_1->RenameSegment(next_seg_id);
    segments_[seg_id] = seg_0;
    segments_[next_seg_id] = seg_1;
    indexes_.insert(indexes_.begin() + seg_index + 1, LinearIndexPair{split_key, next_seg_id});
}

index_seg_id_t LinearIndex::GetNextSegmentId() {
    return indexes_.size();
}

void LinearIndex::HandleSplitSegments() {
    if (need_split_.empty()) {
        return;
    }
    for (int i = need_split_.size() - 1; i >= 0; i--) {
        int seg_index = need_split_[i];
        SplitSegment(seg_index);
    }
    need_split_.clear();
}

std::vector<LinearIndexPair> LinearIndex::GetIndexes() {
    return indexes_;
}

void LinearIndex::QueryIndexes(std::vector<IndexQuery> &queries, std::vector<LinearIndexLeafPair> &res_pairs) {
    LinearIndexKey min_key = LinearIndexKey::GetMinValueOfType(key_type_);
    LinearIndexKey max_key = LinearIndexKey::GetMaxValueOfType(key_type_);
    for (auto &query : queries) {
        if (query.comp == GREATER_THAN || query.comp == GREATER_EQUAL) {
            min_key = query.key;
        } else if (query.comp == LESS_THAN || query.comp == LESS_EQUAL) {
            max_key = query.key;
        } else {
            min_key = max_key = query.key;
        }
    }

    int seg_index = FindFirstSegmentIndex(min_key);
    while (seg_index < indexes_.size()) {
        auto segment = GetSegmentByIndex(seg_index);
        if (segment->GetMinKey().CompareGreaterThan(max_key)) {
            break;
        }
        LinearIndexSegmentReader reader(segment);
        while (reader.HasMoreIndex()) {
            auto leaf_pair = reader.GetNextIndexPair();
            bool valid = true;
            for (auto &query : queries) {
                if (!leaf_pair.key.Compare(query.comp, query.key)) {
                    valid = false;
                    break;
                }
            }
            if (valid) {
                res_pairs.push_back(leaf_pair);
            }
        }
        seg_index++;
    }
}

int LinearIndex::FindFirstSegmentIndex(LinearIndexKey min_key) {
    if (indexes_.size() == 1 || (indexes_.size() > 1 && indexes_[1].key.CompareGreaterThan(min_key))) {
        return 0;
    }
    auto index_info = std::upper_bound(indexes_.begin(), indexes_.end(), min_key,
                                       [](LinearIndexKey key, const LinearIndexPair &index) {
                                           return index.key.CompareGreaterThanEquals(key);
                                       });
    return index_info - indexes_.begin();
}

std::shared_ptr<LinearIndexSegment> LinearIndex::GetSegmentByIndex(int index) {
    LinearIndexPair &pair = indexes_[index];
    index_seg_id_t seg_id = pair.seg_id;
    return GetIndexSegment(seg_id);
}

void LinearIndex::GetAllLeafIndexPairs(std::vector<LinearIndexLeafPair>& res_pairs) {
    for (int seg_index = 0; seg_index < indexes_.size(); seg_index++) {
        auto segment = GetSegmentByIndex(seg_index);
        LinearIndexSegmentReader reader(segment);
        while (reader.HasMoreIndex()) {
            auto leaf_pair = reader.GetNextIndexPair();
            res_pairs.push_back(leaf_pair);
        }
    }
}

}